import java.net.URI;
import java.net.http.WebSocket;
import java.net.http.HttpClient;
import java.net.http.WebSocket.Listener;
import java.util.concurrent.CompletionStage;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;
import java.nio.ByteBuffer;
import java.util.concurrent.atomic.AtomicInteger;

public class WebSocketLoadTest {
    private static final String WS_URL = "ws://localhost:8090/barrage";
    private static final int CLIENT_COUNT = 100; // 并发客户端数
    private static final int MESSAGES_PER_CLIENT = 10; // 每个客户端发送的消息数
    
    private static AtomicInteger connectedClients = new AtomicInteger(0);
    private static AtomicInteger messageCount = new AtomicInteger(0);
    private static CountDownLatch connectLatch = new CountDownLatch(CLIENT_COUNT);
    private static CountDownLatch messageLatch = new CountDownLatch(CLIENT_COUNT * MESSAGES_PER_CLIENT);
    
    public static void main(String[] args) throws Exception {
        System.out.println("开始WebSocket压测...");
        System.out.println("目标URL: " + WS_URL);
        System.out.println("并发客户端数: " + CLIENT_COUNT);
        System.out.println("每个客户端发送消息数: " + MESSAGES_PER_CLIENT);
        
        long startTime = System.currentTimeMillis();
        
        // 创建多个WebSocket客户端
        for (int i = 0; i < CLIENT_COUNT; i++) {
            createWebSocketClient(i);
            Thread.sleep(10); // 稍微延迟，避免连接过于集中
        }
        
        // 等待所有客户端连接完成
        if (connectLatch.await(30, TimeUnit.SECONDS)) {
            System.out.println("所有客户端连接成功: " + connectedClients.get());
        } else {
            System.out.println("连接超时，已连接客户端数: " + (CLIENT_COUNT - connectLatch.getCount()));
        }
        
        // 等待所有消息发送完成
        if (messageLatch.await(60, TimeUnit.SECONDS)) {
            long endTime = System.currentTimeMillis();
            System.out.println("压测完成!");
            System.out.println("总发送消息数: " + messageCount.get());
            System.out.println("总耗时: " + (endTime - startTime) + "ms");
            System.out.println("平均TPS: " + (messageCount.get() * 1000.0 / (endTime - startTime)));
        } else {
            System.out.println("消息发送超时");
        }
    }
    
    private static void createWebSocketClient(int clientId) {
        HttpClient client = HttpClient.newHttpClient();
        WebSocket.Builder builder = client.newWebSocketBuilder();
        
        TestWebSocketListener listener = new TestWebSocketListener(clientId);
        
        builder.buildAsync(URI.create(WS_URL), listener)
               .thenAccept(webSocket -> {
                   connectedClients.incrementAndGet();
                   connectLatch.countDown();
                   
                   // 发送消息
                   for (int i = 0; i < MESSAGES_PER_CLIENT; i++) {
                       String message = String.format(
                           "{\"userId\":\"user%d\",\"username\":\"User%d\",\"content\":\"测试弹幕消息%d\",\"color\":\"#FF0000\",\"fontSize\":16,\"position\":\"normal\"}",
                           clientId, clientId, i
                       );
                       webSocket.sendText(message, true);
                       try {
                           Thread.sleep(100); // 每100ms发送一条消息
                       } catch (InterruptedException e) {
                           Thread.currentThread().interrupt();
                       }
                   }
               })
               .exceptionally(throwable -> {
                   System.err.println("客户端 " + clientId + " 连接失败: " + throwable.getMessage());
                   connectLatch.countDown();
                   return null;
               });
    }
    
    static class TestWebSocketListener implements Listener {
        private int clientId;
        
        public TestWebSocketListener(int clientId) {
            this.clientId = clientId;
        }
        
        @Override
        public void onOpen(WebSocket webSocket) {
            System.out.println("客户端 " + clientId + " 连接成功");
            Listener.super.onOpen(webSocket);
        }
        
        @Override
        public CompletionStage<?> onText(WebSocket webSocket, CharSequence data, boolean last) {
            System.out.println("客户端 " + clientId + " 收到消息: " + data);
            messageCount.incrementAndGet();
            messageLatch.countDown();
            return Listener.super.onText(webSocket, data, last);
        }
        
        @Override
        public void onError(WebSocket webSocket, Throwable error) {
            System.err.println("客户端 " + clientId + " 发生错误: " + error.getMessage());
            connectLatch.countDown();
            messageLatch.countDown();
            Listener.super.onError(webSocket, error);
        }
    }
}